import traceback
from tensorflow.python.keras.preprocessing.image import ImageDataGenerator
import tensorflow as tf
from conf import Conf


class F1_score(tf.keras.metrics.Metric):
    def __init__(self, thresholds=0.5, name='f1_score', **kwargs):
        super(F1_score, self).__init__(name=name, **kwargs)
        self.f1 = self.add_weight(name='f1', initializer='zeros')
        self.tp = self.add_weight(name='tp', initializer='zeros')
        self.fp = self.add_weight(name='fp', initializer='zeros')
        self.fn = self.add_weight(name='fn', initializer='zeros')
        self.thresholds = thresholds

    def update_state(self, y_true, y_pred, sample_weight=None):
        min_delta = 1e-6
        y_pred = tf.cast(tf.where(y_pred > self.thresholds, 1, 0), tf.int8)
        y_true = tf.cast(y_true, tf.int8)

        tp = tf.math.count_nonzero(y_pred * y_true, dtype=tf.float32)
        fp = tf.math.count_nonzero(y_pred * (1 - y_true), dtype=tf.float32)
        fn = tf.math.count_nonzero((1 - y_pred) * y_true, dtype=tf.float32)

        self.tp.assign_add(tp)
        self.fp.assign_add(fp)
        self.fn.assign_add(fn)

        self.f1.assign(2 * self.tp / (2 * self.tp + self.fp + self.fn + min_delta))

    def result(self):
        return self.f1

    def reset_states(self):
        self.f1.assign(0.)
        self.tp.assign(0.)
        self.fp.assign(0.)
        self.fn.assign(0.)


class Util(Conf):


    def get_data(self):
        try:
            datagen = ImageDataGenerator(
                # 归一化处理
                rescale=1. / 255,
                validation_split=self.validation_split
            )
            train_ds = datagen.flow_from_directory(
                self.data_dir,
                seed=123,
                class_mode='categorical',
                target_size=(self.img_height, self.img_width),
                subset='training',
                batch_size=self.batch_size

            )
            val_ds = datagen.flow_from_directory(
                self.data_dir,
                seed=123,
                class_mode='categorical',
                target_size=(self.img_height, self.img_width),
                subset='validation',
                batch_size=self.batch_size)
        except Exception as e:
            traceback.print_exc()
        return train_ds, val_ds
